#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Purpose: Excel export extension
# Created: 02/24/2009
# $Id $
# Copyright (c) 2008,2009 by University College London
# Authors:
# Didrik Pinte --<dpinte@dipole-consulting.com>
#
# This program is free software under the GPL (>=v2)
# Read the file COPYING coming with HydroPlatform for details.
'''
Excel exporters modules based on xlwt
'''

from itertools import count
import logging
import numpy
from sqlalchemy import and_

from hydroplatform.model.export import Exporter, ExportError, ExporterRegistry
from hydroplatform.model import network as nt
from hydroplatform.model import fields as fd

from common import get_field


try:
    import xlwt
    # register exporter 
    class ExcelExporter(Exporter):
        '''
        Exports network to a simple Excel file that should be easily 
        imported in GAMS for example
        '''
        _name = "Excel exporter for GAMS/GDX"
        _description  = '''Exports a network to an Excel file '''
        
        def __init__(self):
            Exporter.__init__(self)
            self.session = None
        
        def export(self, network, filename, session=None):
            '''
            Export an Excel workbook to the given filename
            '''
            if session is None:
                raise ExportError(u"Cannot export without a session object")
            self.session = session
            assert isinstance(network, nt.Network)
            # Create a new workbook to work on *
            workbook = xlwt.Workbook(encoding='cp1251')

            # write the node sheet
            node_sheet = workbook.add_sheet(u"Nodes")             
            ntypes = dict()
            # get node list for this network ordered by node type
            nodes_query = session.query(nt.Node).filter(\
                nt.Node.networks.any(nt.Network.id == network.id))\
                .order_by(nt.Node.type_id)
            nodes = nodes_query.all()
            for i, node in enumerate(nodes):                
                node_sheet.write(i, 0, node.name)
                if not ntypes.has_key(node.type):
                    ntypes[node.type] = [node]
                else: 
                    ntypes[node.type].append(node)
                
            # for each node type, create a sheet with its node list
            for ntype, nodes in ntypes.iteritems():
                ntype_sheet = workbook.add_sheet(ntype.name)
                for i, node in enumerate(nodes):
                    ntype_sheet.write(i, 0, node.name)
                    
            ltypes = dict()
            # links list
            links_sheet = workbook.add_sheet(u"Links")
            for i, link in enumerate(network.links):
                links_sheet.write(i, 0, link.start.name)
                links_sheet.write(i, 1, link.end.name)
                if not ltypes.has_key(link.type):
                    ltypes[link.type] = []
                ltypes[link.type].append(link)
            # for each link type, create a sheet with its link list
            for ltype, links in ltypes.iteritems():
                ltype_sheet = workbook.add_sheet(ltype.name)
                for i, link in enumerate(links):
                    ltype_sheet.write(i, 0, link.start.name)                
                    ltype_sheet.write(i, 1, link.end.name)                
             
            # links connectivity    
            connectivity_sheet = workbook.add_sheet(u"Connect")
            # write the first column headers
            nodes = nodes_query.all()
            for j, node in enumerate(nodes):
                connectivity_sheet.write(0, j+1, node.name)   
            # write the next columns 
            startnodes = session.query(nt.Node).filter(\
                nt.Node.networks.any(nt.Network.id == network.id))\
                .order_by(nt.Node.type_id).all()            
            endnodes = session.query(nt.Node).filter(\
                nt.Node.networks.any(nt.Network.id == network.id))\
                .order_by(nt.Node.type_id).all()

            for i, snode in enumerate(startnodes):
                # the first line is the node name
                connectivity_sheet.write(i+1, 0, snode.name)
                for col, enode in enumerate(endnodes):
                    if network.has_link(snode, enode):
                        for link  in network.links:
                            if link.start == snode and link.end == enode :
                                # not exporting the link.type_id but O/1
                                #connectivity_sheet.write(i+1, col+1, link.type_id)
                                connectivity_sheet.write(i+1, col+1, 1)
                                
            # for each node type
            for ntype, nodes in ntypes.iteritems():
                logging.debug("Node type : %s" % ntype.name)
                self.export_otype(workbook, ntype, nodes)
                
            # for each link type
            for ltype, links in ltypes.iteritems():
                logging.debug("Link type : %s" % ltype.name)
                self.export_otype(workbook, ltype, links)
                    
            
            # Save the file with any name *
            workbook.save(filename) 
        
            
        def export_otype(self, workbook, otype, objects):
            """
            Export all the fields for the given object type into the selected 
            workbook
            """
           # parameter type in a group are homogenous
            if len(otype.fields) == 0:
                return
            
            # group fields by type
            parameters = []
            timeseries = []
            tables = []
            sparams = []
            has_params = 0
            has_timeseries = 0
            has_tables = 0
            has_sparam = 0

            for field in otype.fields:
                logging.debug('--> %s' % field.name)
                if isinstance(field.attribute_type, fd.Parameter):
                    parameters.append(field)
                    has_params = 1
                elif isinstance(field.attribute_type, fd.TimeSerie):
                    timeseries.append(field)
                    has_timeseries = 1
                elif isinstance(field.attribute_type, fd.Table):
                    tables.append(field)
                    has_tables = 1
                elif isinstance(field.attribute_type, fd.SeasonalParameter):
                    sparams.append(field)
                    has_sparam = 1
             
            # add a sheet per field types per type if needed
            
            if len(parameters) > 0 :
                name = self.get_sheet_name(otype, None, "par")
                field_sheet = workbook.add_sheet(name)
                self.export_parameters(field_sheet, 
                                             objects, 
                                             parameters)

            if len(timeseries) > 0:
                name = self.get_sheet_name(otype, None, "ts")
                field_sheet = workbook.add_sheet(name)
                self.export_timeseries(field_sheet, 
                                             objects, 
                                             timeseries)
                
            if len(tables) > 0:
                # Table export is one table per sheet per node type
                for table in tables:
                    name = u"%s_tab_%s" % (otype.name.upper(),
                                       table.name)
                    table_sheet = workbook.add_sheet(name)
                    self.export_table(table_sheet, objects, table)
                    
            if len(sparams) > 0:
                # Table export is one table per sheet per node type
                for sparam in sparams:
                    name = u"%s_%s" % (otype.name.upper(),
                                       sparam.name)
                    logging.debug("SP param sheet : " + name)
                    table_sheet = workbook.add_sheet(name)
                    self.export_seasonal_parameter(table_sheet, objects, sparam)
            
        def export_group(self, workbook, otype, group, objects):
            '''
            Export fields for the given group into the selected workbook and 
            for the given object type
            '''
            if isinstance(otype, nt.NodeType):
                fields = group.nodefields
            elif isinstance(otype, nt.LinkType):
                fields = group.linkfields 
                
           # parameter type in a group are homogenous
            if len(fields) == 0:
                return
            
            # group fields by type
            parameters = []
            timeseries = []
            tables = []
            sparams = []
            has_params = 0
            has_timeseries = 0
            has_tables = 0
            has_sparam = 0

            for field in fields:
                # filter only fields that belongs to this nodetype
                if field.type is not otype:
                    continue
                logging.debug('--> %s' % field.name)
                if isinstance(field.attribute_type, fd.Parameter):
                    parameters.append(field)
                    has_params = 1
                elif isinstance(field.attribute_type, fd.TimeSerie):
                    timeseries.append(field)
                    has_timeseries = 1
                elif isinstance(field.attribute_type, fd.Table):
                    tables.append(field)
                    has_tables = 1
                elif isinstance(field.attribute_type, fd.SeasonalParameter):
                    sparams.append(field)
                    has_sparam = 1
             
            # if the group has multiple types, the name of the sheet that will
            # be created of each type for the group must be distinct
            # BUT if node has only one group, the name of the sheet will not 
            # have the group name
            has_multiple_types = has_params + has_timeseries + has_tables \
                               + has_sparam
            
            # add a sheet per group per type if needed
            
            if len(parameters) > 0 :
                name = self.get_sheet_name(otype, group, "par", \
                                          has_multiple_types)
                group_sheet = workbook.add_sheet(name)
                self.export_parameters(group_sheet, 
                                             objects, 
                                             parameters)

            if len(timeseries) > 0:
                name = self.get_sheet_name(otype, group, "ts",\
                                          has_multiple_types)
                group_sheet = workbook.add_sheet(name)
                self.export_timeseries(group_sheet, 
                                             objects, 
                                             timeseries)
                
            if len(tables) > 0:
                # Table export is one table per sheet per node type
                for table in tables:
                    name = u"%s_%s" % (otype.name.upper(),
                                       table.name)
                    table_sheet = workbook.add_sheet(name)
                    self.export_table(table_sheet, objects, table)
                    
            if len(sparams) > 0:
                # Table export is one table per sheet per node type
                for sparam in sparams:
                    name = u"%s_%s" % (otype.name.upper(),
                                       sparam.name)
                    table_sheet = workbook.add_sheet(name)
                    self.export_seasonal_parameter(table_sheet, objects, sparam)

        def get_sheet_name(self, nodetype, group, ftype, has_multiple_types=False):
            
            if group is not None and len(nodetype.groups) > 1:
                name = (has_multiple_types > 1) \
                     and u"%%s_%%s_%s" % ftype \
                     or u"%s_%s"
                return name % (nodetype.name.upper(), 
                               group.name.lower())
            else:
                name = u"%%s_%s" % ftype
                return name % nodetype.name.upper()            
 
                
        def export_parameters(self, sheet, objects, parameters):
            '''
            Group parameters are outputted like this :
            
                    Param1  Param2
            Node1    v1       v2
            Node2    v3       v4
            '''
            nobject = objects[0]
            # get value for the given field
            if isinstance(nobject, nt.Node):
                dbtable = nt.NodeValue
                key_id = dbtable.node_id
                col_idx = 1
            elif isinstance(nobject, nt.Link):
                dbtable = nt.LinkValue
                key_id = nt.LinkValue.link_id
                col_idx = 2
                
            # write field name in column
            for col, field in enumerate(parameters):
                sheet.write(0, col+col_idx, field.name) 
                
            # for each node write its name and values for the each field
            obj_count = count()
            row = obj_count.next()
            for nobject in objects:
                node_has_value = False
                for col, field in enumerate(parameters):
                    val = get_field(nobject, field, self.session)
                    if val is not None and val.value is not None:          
                        # output data    
                        node_has_value = True
                        sheet.write(row+1, col+col_idx, unicode(val.value))
                if node_has_value is True:
                    if isinstance(nobject, nt.Node):
                        sheet.write(row+1, 0, nobject.name)
                    elif isinstance(nobject, nt.Link):
                        sheet.write(row+1, 0, nobject.start.name)
                        sheet.write(row+1, 1, nobject.end.name)
                    row = obj_count.next()
         
 
        def export_seasonal_parameter(self, sheet, nodes, field):
            """
            Exports a table to a sheet 
            
            Group table are outputted like this :
            date     node1-value   node2-value ... 
            Y-m-d       v1               v2   
            Y-m-d       v1               v2   
            Y-m-d       v1               v2   
            Y-m-d       v1               v2   
            Y-m-d       v1               v2   
 
            """
            # write column labels
            sheet.write(0, 0, u"Season")
            for i,date in enumerate(field.attribute_type._dates[:-1]):
                sheet.write(1+i, 0, i+1)
            for node_idx, node in enumerate(nodes):
                node_has_value = False
                # get the value field
                val = self.session.query(nt.NodeValue).filter( \
                    and_(nt.NodeValue.attribute_id ==field.id, \
                         nt.NodeValue.node_id == node.id)).first()
                if val is not None:                            
                    # loop on the rows and write the lines
                    matrix = val.value
                    if matrix is None or \
                       numpy.all(numpy.isnan(matrix)):
                        continue
                    # write node name
                    sheet.write(0, 1+node_idx , node.name) 
                    for row in xrange(matrix.shape[0]-1):
                            if not numpy.isnan(matrix[row, 0]) :
                                sheet.write(1+row, 1+node_idx, \
                                        str(matrix[row, 0]))
                            else:
                                continue

                    
        def export_table(self, sheet, nodes, field):
            """
            Exports a table to a sheet 
            
            Group table are outputted like this :
                         tab1-collabel1    tab1-collabel2    tab1-collabel3
            node1  row1     v1               v2                  v3
            node1  row2     vx               v9                  vz
            node1  row3     v1               v2                  v3
            node2  row1     v1               v2                  v3
            node2  row2     vx               v9                  vz
            node2  row3     v1               v2                  v3
 
            """
            # write column labels
            for i, col in enumerate(field.attribute_type.columns):
                sheet.write(0, 2+i, col.name)
            row_count = count()
            rc = row_count.next()
            for node in nodes:
                # get the value field
                val = self.session.query(nt.NodeValue).filter( \
                    and_(nt.NodeValue.attribute_id ==field.id, \
                         nt.NodeValue.node_id == node.id)).first()
                if val is not None:                            
                    # loop on the rows and write the lines
                    matrix = val.value
                    if matrix is None or \
                       numpy.all(numpy.isnan(matrix)):
                        continue
                    for row in xrange(matrix.shape[0]):
                        # write node name
                        sheet.write(1+rc, 0 , node.name)
                        # add dummy variable for GAMS
                        sheet.write(1+rc, 1, u"step%s" % row)
                        for col in xrange(matrix.shape[1]):
                            if not numpy.isnan(matrix[row, col]) :
                                sheet.write(1+rc, 2+col, \
                                        str(matrix[row, col]))
                            else:
                                continue
                        rc = row_count.next()
                        
 
                    
        def export_group_tables(self, sheet, nodes, tables):
            '''
            Group table are outputted like this :
                         tab1-collabel1    tab1-collabel2    tab1-collabel3
            node1  row1     v1               v2                  v3
            node1  row2     vx               v9                  vz
            node1  row3     v1               v2                  v3
            node2  row1     v1               v2                  v3
            node2  row2     vx               v9                  vz
            node2  row3     v1               v2                  v3
            
            Multiple tables are concatenated on the right
            '''
            colidx = 1
            rowidx = 1
            # for each table --> write the table
            for field in tables:
                # write column labels
                for i, label in enumerate(field.attribute_type.colLabels):
                    sheet.write(0, colidx + i, label)
                for node in nodes:
                    node_has_value = False
                    # get the value field
                    val = self.session.query(nt.NodeValue).filter( \
                        and_(nt.NodeValue.attribute_id ==field.id, \
                             nt.NodeValue.node_id == node.id)).first()
                    if val is not None:                            
                        # loop on the rows and write the lines
                        matrix = val.value
                        if matrix is None or \
                           numpy.all(numpy.isnan(matrix)):
                            continue
                        for row in xrange(matrix.shape[0]):
                            # write node name
                            sheet.write(rowidx + row, colidx-1, node.name)
                            for col in xrange(matrix.shape[1]):
                                if not numpy.isnan(matrix[row, col]) :
                                    sheet.write(rowidx + row, colidx + col, \
                                            str(matrix[row, col]))
                                else:
                                    continue
                        rowidx += matrix.shape[0]
                # update the column index for the next table
                colidx = len(field.attribute_type.colLabels) + 1
                
        def export_timeseries(self, sheet, nodes, timeseries):
            '''
            Group timeseries are outputted like this :
                       TS1    TS1      TS2
                       Node1  Node2   Node1
            datetime1   ...    ...      ...
            datetime2   ...    ...      ...
            datetime3
            
            Multiple timeseries are concatenated on the right
            '''
            has_row_label = False
            skipped_node = 0

            
            if isinstance(nodes[0], nt.Node):
                header_count = 2
            elif isinstance(nodes[0], nt.Link):
                header_count = 3
                
            col_count = count(3)
            col = col_count.next()           
            for i, field in enumerate(timeseries):
                for j, node in enumerate(nodes):
                    node_has_value = False
                    # get ts for field
                    val = get_field(node, field, self.session)
                    if val is not None:                            
                        matrix = val.value
                        if matrix is None or \
                           numpy.all(numpy.isnan(matrix)) is True:
                            continue
                        else:
                            node_has_value = True
                        if has_row_label is False:
                            sheet.write(0, 0, u"Year")
                            sheet.write(0, 1, u"Month")
                            sheet.write(0, 2, u"Day")
                            # write row labels
                            for row, timestamp in enumerate(field.attribute_type.dates):
                                sheet.write(row+header_count, 0, timestamp.year)
                                sheet.write(row+header_count, 1, timestamp.month)
                                sheet.write(row+header_count, 2, timestamp.day)
                            has_row_label = True
                        # write data
                        for row in xrange(matrix.shape[0]):
                            sheet.write(row+header_count, col, str(matrix[row, 0]))
                    # write header
                    if node_has_value:
                        sheet.write(0, col, field.name)
                        if header_count is 2:
                            sheet.write(1, col, node.name)
                        else:
                            # this is a link
                            sheet.write(1, col, node.start.name)
                            sheet.write(2, col, node.end.name)
                            
                        col = col_count.next()
 
    REGISTRY = ExporterRegistry()
    REGISTRY.register(ExcelExporter)
    logging.info("ExcelExporter extension registered")
    
except ImportError:
    logging.warning('XLWT not available - Excel export extension disabled')